// ragflow-upload.js - 直接上传 .md → .txt（不清洗、不处理）

import fs from 'fs';
import path from 'path';
import axios from 'axios';
import FormData from 'form-data';
import process from 'process';

// ✅ 配置区
const RAGFLOW_URL = 'http://rf-dify.docker.internal:5175';
const KB_ID = '3dd5d3bcb0d711f0abbd061f9ce6bf51';
const AUTH_TOKEN = 'IjA3YzUwMTQ4YjU0NzExZjBhN2RkZmU1NmE1MDA4MTdiIg.aQLmUA.cDBOwfJ2_VNk906IQ6wYDnrcEVA'; // 替换
const METADATA_FILE = './docs/知源中医/知源中医_site.json';
const BOOKS_DIR = './docs/知源中医';
const PAGE_SIZE = 1000;
const headers = { 'Authorization': AUTH_TOKEN, 'Content-Type': 'application/json;charset=UTF-8' }

// 读取元数据
const metadataList = JSON.parse(fs.readFileSync(METADATA_FILE, 'utf-8'));
const metadataMap = new Map();
metadataList.forEach(book => {
  metadataMap.set(book.name, book);
});

// 获取知识库中所有文档列表
async function listDocuments() {
  try {
    const response = await axios.post(`${RAGFLOW_URL}/v1/document/list`, { kb_id: KB_ID, page: 1, page_size: 1000 }, { headers });
    console.log(response)
    return response.data.data;
  } catch (error) {
    console.error('❌ 获取文档列表失败:', error.response?.data);
    return [];
  }
}

// 删除文档（你已有）
async function deleteDocument(docId) {
  try {
    const response = await axios.post(`${RAGFLOW_URL}/v1/document/rm`, { doc_id: docId }, { headers });
    console.log(`🗑️ 已删除文档: ${docId}`);
    if (response.data.code != 0) {
      console.log(`${RAGFLOW_URL}/v1/document/rm`, { doc_id: docId }, { headers }, response.data)
      process.exit(1)
    }
  } catch (error) {
    console.error(`❌ 删除失败: ${docId}`, error.response?.data);
  }
}

// 批量清理旧的 .md 文档
async function cleanupOldMdDocuments() {
  console.log('🔍 正在查找旧的 .md 文档...');

  const docs = await listDocuments();
  const mdDocs = docs.filter(doc => doc.name.endsWith('.md'));

  if (mdDocs.length === 0) {
    console.log('✅ 未发现需要清理的 .md 文档');
    return;
  }

  console.log(`⚠️  发现 ${mdDocs.length} 个旧的 .md 文档，准备删除:`);
  for (const doc of mdDocs) {
    console.log(`   📄 ${doc.name} (ID: ${doc.id})`);
    await deleteDocument(doc.id);
  }

  console.log('✅ 清理完成！');
}

// 查询列表
async function getDocsList(rs) {
  const url = `${RAGFLOW_URL}/v1/document/list`;
  const result = new Map();

  console.log(`🔍 查询未解析的 .txt 文件（page_size=${PAGE_SIZE}）...`);
  let response = null;
  try {
    let page = 1;
    while (true) {
      response = await axios.post(`${url}?kb_id=${KB_ID}&page_size=${PAGE_SIZE}&page=${page}`, { suffix: ['txt'], run_status: rs }, { headers });
      // console.log(`${url}?kb_id=${KB_ID}&page_size=${PAGE_SIZE}&page=${page}`, { suffix: ['txt'], run_status: rs }, { headers })
      const { docs, total } = response.data.data;
      if (!docs || docs.length === 0) break;

      docs.forEach(doc => {
        const decodedName = decodeUnicode(doc.name);
        result.set(decodedName, doc.id);
      });

      console.log(`📥 第 ${page} 页: 获取 ${docs.length} 条，累计 ${result.size}/${total}`);
      if (result.size >= total) break;
      page++;
    }

    console.log(`✅ 共找到 ${result.size} 个未解析的 .txt 文件`);
  } catch (error) {
    console.error('❌ 查询失败:', error);
    console.log(response.data)
    process.exit()
  }

  return result;

  function decodeUnicode(str) {
    try {
      return JSON.parse(`"${str}"`);
    } catch (e) {
      return str;
    }
  }
}

// 检查文件是否已存在
async function documentExists(compositeName) {
  const keywords = encodeURIComponent(compositeName);
  const url = `${RAGFLOW_URL}/v1/document/list?kb_id=${KB_ID}&keywords=${keywords}&page_size=50&page=1`;

  try {
    const response = await axios.get(url, { headers });
    const documents = response.data.data?.list || [];
    return documents.some(doc => doc.name === compositeName);
  } catch (error) {
    console.error(`🔍 查询失败 (${compositeName}):`, error.message);
    return false;
  }
}

// 上传文件（.md → 读取内容 → 内存中转.txt → 上传）
async function uploadFile(filePath, filename, bookInfo) {
  const compositeName = `${bookInfo.name}__${filename.replace('.md', '.txt')}`;

  try {
    const content = fs.readFileSync(filePath, 'utf-8');

    // ✅ 使用 form-data 替代原生 FormData
    const form = new FormData();

    // ✅ 直接 append 字符串内容，指定 filename 和 contentType
    form.append('file', content, {
      filename: compositeName,
      contentType: 'text/plain',
    });
    form.append('kb_id', KB_ID);
    form.append('parser_id', 'book');

    const response = await axios({
      method: 'post',
      url: `${RAGFLOW_URL}/v1/document/upload`,
      data: form,
      headers: {
        ...form.getHeaders(), // ✅ form-data 提供的 headers
        'Authorization': AUTH_TOKEN,
      },
      timeout: 30000
    });
    if (response.data.code != 0) {
      console.log(`${RAGFLOW_URL}/v1/document/upload`, form, {
        headers: {
          ...form.getHeaders(), // ✅ form-data 提供的 headers
          ...headers
        }
      }, response.data)
      process.exit(1)
    }
    console.log(`✅ 成功上传: ${compositeName}`);
    return response.data.data[0]?.id;
  } catch (error) {
    const msg = error.response?.data?.message || error.message;
    console.error(`❌ 上传失败 (${compositeName}):`, msg);
    return null;
  }
}

// 触发解析
async function triggerParse(docId) {
  try {
    await axios.post(`${RAGFLOW_URL}/v1/document/run`, { doc_ids: [docId], run: 1, delete: false }, { headers });
    console.log(`▶️  已触发解析: ${docId}`);
  } catch (error) {
    console.error(`❌ 触发解析失败 (${docId}):`, error.response?.data || error.message);
  }
}

// 主函数
async function main() {
  // cleanupOldMdDocuments();
  // return
  const isParse = process.argv.includes('--parse');
  const isDelete = process.argv.includes('--delete');
  console.log(`🚀 开始上传${isParse ? '并解析' : ''}...`);

  const entries = fs.readdirSync(BOOKS_DIR, { withFileTypes: true });

  if (isDelete) {
    const allNames = await getDocsList([]);
    for (let [compositeName, docId] of allNames) {
      console.log(`   📄 ${compositeName} (ID: ${docId})`);
      await deleteDocument(docId);  // 会按顺序等待每个删除完成
    }
    console.log('🎉 删除任务完成！');
    return
  }

  const unparsedNames = await getDocsList([]);
  for (const entry of entries) {
    if (entry.isDirectory()) {
      const bookInfo = metadataMap.get(entry.name);
      if (!bookInfo) continue;

      const bookPath = path.join(BOOKS_DIR, entry.name);
      const files = fs.readdirSync(bookPath).filter(f => f.endsWith('.md'));

      for (const file of files) {
        const compositeName = `${bookInfo.name}__${file.replace('.md', '.txt')}`;

        if (unparsedNames.has(compositeName)) {
          if (isParse) {
            console.log(`⏭️  已存在，解析: ${compositeName}`);
            const docId = unparsedNames.get(compositeName);
            await triggerParse(docId);
          } else {
            console.log(`⏭️  已存在且未解析，跳过: ${compositeName}`);
          }
          continue;
        }

        if (await documentExists(compositeName)) {
          console.log(`⏭️  已存在，跳过: ${compositeName}`);
          continue;
        }

        const docId = await uploadFile(path.join(bookPath, file), file, bookInfo);
        if (docId && isParse) {
          await triggerParse(docId);
        }
      }
    }
  }

  console.log('🎉 上传任务完成！');
}

// 启动
main();